#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
使用SerpApi获取文章引用信息的脚本
每次运行处理data/source_article.txt中的1篇文章，
获取引用文章信息并直接保存为source_article,url格式
"""

import os
import sys
import time
import logging
from typing import List, Dict, Optional, Tuple
from datetime import datetime
from pathlib import Path

# 添加项目路径到sys.path
project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))  # 回到crawl_by_trae目录
sys.path.insert(0, project_root)

try:
    from serpapi.fetch import SerpApiFetcher
    from serpapi.config import TARGET_YEAR
except ImportError as e:
    print(f"导入错误: {e}")
    print("请确保crawl_by_trae/serpapi模块可用")
    sys.exit(1)

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('run_fetch.log', encoding='utf-8'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)


class ArticleFetcher:
    """文章引用信息获取器 - 支持用户指定行号"""
    
    def __init__(self):
        """初始化"""
        self.fetcher = SerpApiFetcher()
        self.source_file = "data/source_article.txt"
        self.output_dir = "data"
        self.formatted_output_file = "data/citing_articles_formatted01.txt"  # 格式化输出文件（带编号01、02、03格式）
        
        # 确保输出目录存在
        os.makedirs(self.output_dir, exist_ok=True)
        
        logger.info("ArticleFetcher初始化完成")
    
    def generate_output_filename(self, line_numbers: List[int]) -> str:
        """
        根据行号范围生成输出文件名
        
        Args:
            line_numbers: 处理的行号列表
            
        Returns:
            生成的文件名
        """
        if not line_numbers:
            return "data/citing_articles_formatted.txt"
        
        # 排序行号列表
        sorted_lines = sorted(line_numbers)
        min_line = sorted_lines[0]
        max_line = sorted_lines[-1]
        
        # 生成文件名格式：citing_articles_formatted_1-10.txt
        if min_line == max_line:
            filename = f"data/citing_articles_formatted_{min_line}.txt"
        else:
            filename = f"data/citing_articles_formatted_{min_line}-{max_line}.txt"
        
        return filename
    
    def parse_line_input(self, user_input: str) -> List[int]:
        """
        解析用户输入的行号
        
        支持的格式:
        - 单个行号: "5"
        - 多个行号: "1,3,5"
        - 范围: "1-5" 或 "1:5"
        - 混合: "1,3-5,8"
        
        Args:
            user_input: 用户输入的字符串
            
        Returns:
            行号列表（1-based）
        """
        line_numbers = []
        
        try:
            # 分割逗号分隔的部分
            parts = user_input.strip().split(',')
            
            for part in parts:
                part = part.strip()
                
                # 检查是否是范围（支持 - 和 : 两种分隔符）
                if '-' in part:
                    start, end = part.split('-', 1)
                    start, end = int(start.strip()), int(end.strip())
                    line_numbers.extend(range(start, end + 1))
                elif ':' in part:
                    start, end = part.split(':', 1)
                    start, end = int(start.strip()), int(end.strip())
                    line_numbers.extend(range(start, end + 1))
                else:
                    # 单个行号
                    line_numbers.append(int(part))
            
            # 去重并排序
            line_numbers = sorted(list(set(line_numbers)))
            
            # 验证行号有效性
            total_articles = self.get_total_articles()
            valid_lines = [line for line in line_numbers if 1 <= line <= total_articles]
            
            if len(valid_lines) != len(line_numbers):
                invalid_lines = [line for line in line_numbers if line not in valid_lines]
                logger.warning(f"无效的行号将被忽略: {invalid_lines}")
                logger.info(f"有效行号范围: 1-{total_articles}")
            
            return valid_lines
            
        except ValueError as e:
            logger.error(f"行号格式错误: {e}")
            return []
        except Exception as e:
            logger.error(f"解析行号失败: {e}")
            return []
    
    def get_total_articles(self) -> int:
        """获取总文章数量"""
        try:
            if not os.path.exists(self.source_file):
                return 0
            
            with open(self.source_file, 'r', encoding='utf-8') as f:
                lines = f.readlines()
            
            return len([line for line in lines if line.strip()])
            
        except Exception as e:
            logger.error(f"获取文章总数失败: {e}")
            return 0
    
    def get_articles_by_lines(self, line_numbers: List[int]) -> List[Tuple[str, int]]:
        """
        根据行号获取文章列表
        
        Args:
            line_numbers: 行号列表（1-based）
            
        Returns:
            (article_title, line_number) 的列表
        """
        try:
            with open(self.source_file, 'r', encoding='utf-8') as f:
                lines = f.readlines()
            
            # 过滤空行
            valid_lines = [line.strip() for line in lines if line.strip()]
            
            articles = []
            for line_num in line_numbers:
                if 1 <= line_num <= len(valid_lines):
                    articles.append((valid_lines[line_num - 1], line_num))
                else:
                    logger.warning(f"行号 {line_num} 超出范围 (1-{len(valid_lines)})")
            
            return articles
            
        except Exception as e:
            logger.error(f"读取文章失败: {e}")
            return []
    

    
    def process_article(self, article_title: str) -> List[str]:
        """
        处理单篇文章，获取其引用文章的URL列表
        
        Args:
            article_title: 文章标题
            
        Returns:
            引用文章URL列表
        """
        logger.info(f"处理文章: {article_title[:50]}...")
        
        try:
            # 获取文章的cites_id
            cites_id = self.fetcher.get_article_cites_id(article_title)
            if not cites_id:
                logger.warning(f"未找到文章的cites_id: {article_title[:50]}...")
                return []
            
            logger.info(f"找到cites_id: {cites_id}")
            
            # 获取引用文章URL列表
            citing_urls = self.fetcher.get_citing_articles(cites_id, year=TARGET_YEAR, max_results=500)
            
            if not citing_urls:
                logger.info(f"未找到{TARGET_YEAR}年的引用文章")
                return []
            
            logger.info(f"找到 {len(citing_urls)} 篇{TARGET_YEAR}年引用文章")
            return citing_urls
            
        except Exception as e:
            logger.error(f"处理文章失败: {e}")
            return []
    
    def save_formatted_results(self, source_article: str, line_number: int, citing_urls: List[str], 
                             output_filename: str = None, append: bool = True):
        """
        将结果保存为 source_article,编号,url 格式
        
        Args:
            source_article: 源文章标题
            line_number: 文章在源文件中的行号
            citing_urls: 引用文章URL列表
            output_filename: 输出文件名（如果为None则使用默认文件名）
            append: 是否追加到文件（True）还是覆盖文件（False）
        """
        try:
            # 使用指定的文件名或默认文件名
            filename = output_filename or self.formatted_output_file
            
            # 确保输出目录存在
            output_path = Path(filename)
            output_path.parent.mkdir(parents=True, exist_ok=True)
            
            mode = 'a' if append else 'w'
            with open(filename, mode, encoding='utf-8') as f:
                # 清理源文章标题中的逗号，避免CSV格式问题
                clean_source = source_article.replace(',', ';')
                
                if citing_urls:
                    for i, url in enumerate(citing_urls, 1):
                        f.write(f"{clean_source},{str(i).zfill(2)},{url}\n")
                else:
                    # 即使没有引用文章，也记录一条记录表示已处理
                    f.write(f"{clean_source},00,无引用文章\n")
                
            logger.info(f"已将第 {line_number} 行文章的 {len(citing_urls)} 条记录保存到 {filename}")
            
        except Exception as e:
            logger.error(f"保存格式化结果失败: {e}")
    
    def process_articles_by_lines(self, line_numbers: List[int]) -> bool:
        """
        处理指定行号的文章
        
        Args:
            line_numbers: 要处理的行号列表
            
        Returns:
            是否成功处理
        """
        try:
            # 获取指定行号的文章
            articles = self.get_articles_by_lines(line_numbers)
            
            if not articles:
                logger.warning("没有找到有效的文章")
                return False
            
            # 生成动态文件名
            output_filename = self.generate_output_filename(line_numbers)
            logger.info(f"输出文件名: {output_filename}")
            
            logger.info(f"开始处理 {len(articles)} 篇文章...")
            
            success_count = 0
            for i, (article_title, line_number) in enumerate(articles, 1):
                logger.info(f"处理第 {i}/{len(articles)} 篇文章 (行号 {line_number}): {article_title[:50]}...")
                
                # 处理文章
                citing_urls = self.process_article(article_title)
                
                if citing_urls:
                    # 保存结果（第一次写入时覆盖文件，后续追加）
                    append_mode = i > 1
                    self.save_formatted_results(article_title, line_number, citing_urls, 
                                              output_filename, append_mode)
                    success_count += 1
                    logger.info(f"第 {line_number} 行文章处理完成，找到 {len(citing_urls)} 篇引用文章")
                else:
                    # 即使没有引用文章也要记录
                    append_mode = i > 1
                    self.save_formatted_results(article_title, line_number, [], 
                                              output_filename, append_mode)
                    logger.warning(f"第 {line_number} 行文章未找到引用文章")
            
            logger.info(f"批量处理完成！成功处理 {success_count}/{len(articles)} 篇文章")
            return success_count > 0
            
        except Exception as e:
            logger.error(f"批量处理失败: {e}")
            return False
    



def main():
    """主函数"""
    try:
        # 初始化获取器
        fetcher = ArticleFetcher()
        
        # 测试API连接
        if not fetcher.fetcher.test_connection():
            logger.error("API连接测试失败，请检查配置")
            return
        
        # 获取总文章数
        total_articles = fetcher.get_total_articles()
        
        print(f"\n=== 文章引用信息获取工具 ===")
        print(f"总文章数: {total_articles}")
        print(f"输出文件将根据处理的行号范围动态命名")
        
        # 显示前10篇文章作为参考
        print(f"\n前10篇文章预览:")
        try:
            with open(fetcher.source_file, 'r', encoding='utf-8') as f:
                lines = f.readlines()
            valid_lines = [line.strip() for line in lines if line.strip()]
            
            for i, line in enumerate(valid_lines[:10], 1):
                print(f"  {i:3d}. {line[:80]}{'...' if len(line) > 80 else ''}")
            
            if len(valid_lines) > 10:
                print(f"  ... (还有 {len(valid_lines) - 10} 篇文章)")
                
        except Exception as e:
            logger.error(f"读取文章列表失败: {e}")
            return
        
        # 获取用户输入
        print(f"\n请输入要处理的行号:")
        print(f"支持格式:")
        print(f"  - 单个行号: 5")
        print(f"  - 多个行号: 1,3,5")
        print(f"  - 范围: 1-5 或 1:5")
        print(f"  - 混合: 1,3-5,8")
        print(f"  - 输入 'q' 或 'quit' 退出")
        
        while True:
            try:
                user_input = input(f"\n请输入行号 (1-{total_articles}): ").strip()
                
                if user_input.lower() in ['q', 'quit', 'exit']:
                    print("退出程序")
                    return
                
                if not user_input:
                    print("请输入有效的行号")
                    continue
                
                # 解析用户输入
                line_numbers = fetcher.parse_line_input(user_input)
                
                if not line_numbers:
                    print("无效的行号格式，请重新输入")
                    continue
                
                # 显示将要处理的文章
                articles = fetcher.get_articles_by_lines(line_numbers)
                output_filename = fetcher.generate_output_filename(line_numbers)
                
                print(f"\n将要处理 {len(articles)} 篇文章:")
                for article_title, line_num in articles:
                    print(f"  行号 {line_num:3d}: {article_title[:80]}{'...' if len(article_title) > 80 else ''}")
                print(f"\n输出文件: {output_filename}")
                
                # 确认处理
                confirm = input(f"\n确认处理这些文章? (y/n): ").strip().lower()
                if confirm not in ['y', 'yes', '是']:
                    print("取消处理")
                    continue
                
                # 开始处理
                print(f"\n开始处理...")
                success = fetcher.process_articles_by_lines(line_numbers)
                
                if success:
                    print(f"\n处理完成！结果已保存到: {output_filename}")
                else:
                    print(f"\n处理失败或未找到引用文章")
                
                # 询问是否继续
                continue_choice = input(f"\n是否继续处理其他文章? (y/n): ").strip().lower()
                if continue_choice not in ['y', 'yes', '是']:
                    break
                    
            except KeyboardInterrupt:
                print(f"\n\n用户中断处理")
                break
            except Exception as e:
                logger.error(f"处理过程中出错: {e}")
                print(f"出现错误: {e}")
                continue
        
        print(f"\n程序结束")
        
    except KeyboardInterrupt:
        print(f"\n\n用户中断处理")
    except Exception as e:
        logger.error(f"程序执行出错: {e}")
        print(f"程序执行出错: {e}")


if __name__ == "__main__":
    main()